from surprise import SVD, KNNBasic, Dataset, Reader
from surprise import accuracy
from surprise.model_selection import train_test_split  # 导入新的拆分方法
import pandas as pd
import pickle
import os


def main():
    # 1. 从处理后的数据文件读取数据
    try:
        processed_ratings = pd.read_csv('filtered_data/ratings_filtered.csv')
    except FileNotFoundError:
        print("错误：未找到processed_ratings.csv文件，请先运行数据处理脚本生成该文件。")
        return

    # 2. 数据准备：提取Surprise库所需的三列（userId, movieId, rating）
    rating_data = processed_ratings[['USER_MD5', 'MOVIE_ID', 'RATING']].dropna()
    if rating_data.empty:
        print("错误：处理后的评分数据为空，请检查原始数据是否有效。")
        return

    # 3. 数据拆分：80%训练集，20%测试集
    reader = Reader(rating_scale=(1, 5))
    data = Dataset.load_from_df(rating_data, reader)

    # 修复：使用train_test_split方法拆分数据集
    trainset, testset = train_test_split(data, test_size=0.2, random_state=42)

    # 4. 训练SVD矩阵分解模型
    print("开始训练SVD模型...")
    svd_model = SVD()
    svd_model.fit(trainset)
    print("SVD模型训练完成")

# 5. 训练基于物品协同过滤（KNN）的推荐模型
    print("\n开始训练KNN物品协同过滤模型...")
    sim_options = {'name': 'cosine', 'user_based': False}
    knn_model = KNNBasic(sim_options=sim_options)
    knn_model.fit(trainset)
    print("KNN物品协同过滤模型训练完成")


    # 6. 保存模型到本地
    print("\n正在保存模型到本地...")
    print(f"模型保存路径：{os.getcwd()}")  # 显示保存路径
    try:
        with open('svd_recommender_model.pkl', 'wb') as f:
            pickle.dump(svd_model, f)

        with open('knn_recommender_model.pkl', 'wb') as f:
            pickle.dump(knn_model, f)
        print("模型保存成功！生成文件：")
        print("-svd_recommender_model.pkl")
        print("-knn_recommender_model.pkl")
    except Exception as e:
        print(f"模型保存失败：{str(e)}")
        return

    # 7. 在测试集上评估模型
    print("\n开始评估模型性能...")
    svd_predictions = svd_model.test(testset)
    knn_predictions = knn_model.test(testset)

    print(f"SVD模型RMSE: {accuracy.rmse(svd_predictions):.4f}")
    print(f"KNN模型RMSE: {accuracy.rmse(knn_predictions):.4f}")
    



if __name__ == "__main__":
    main()